import { GetMap, HasUrl } from "@/decorator/index";
import { Middleware, IMiddleware } from "@midwayjs/core";
import { NextFunction, Context } from "@midwayjs/koa";
import * as ratelimit from "koa-ratelimit"

@Middleware()
export class RateLimitMiddleware implements IMiddleware<Context, NextFunction> {
    resolve(): any {
        return ratelimit({
            driver: "memory",
            db: new Map(), // 使用内存存储方式存储计数器状态
            duration: (ctx: Context) => {
                return (GetMap().get(HasUrl(ctx))?.duration ?? 1) * 1000
            }, // 限制时间窗口，这里是60秒
            max: (ctx: Context) => {
                return GetMap().get(HasUrl(ctx))?.max ?? 10
            }, // 在时间窗口内允许的最大请求数
            id: (ctx: Context) => {
                return GetMap().get(HasUrl(ctx))?.type === "1" ? GetMap().get(HasUrl(ctx))?.url : ctx.ip + (HasUrl(ctx) ?? ctx.path)
            }, // 根据IP地址标识用户
            disableHeader: false, // 不禁用RateLimit-*头
            throw: true, // 超出限制时抛出错误
            headers: {
                remaining: "Rate-Limit-Remaining",
                reset: "Rate-Limit-Reset",
                total: "Rate-Limit-Total"
            },
            // 自定义错误响应
            errorMessage: "您已超过请求限制，请稍后重试！",
        });
    }
}
